package com.augrain.microframework.limit.aspect;

import com.augrain.microframework.limit.LimitKeyGenerator;
import com.augrain.microframework.limit.annotation.RateLimiter;
import com.augrain.microframework.limit.annotation.RateRule;
import com.augrain.microframework.limit.breaker.CircuitBreakerManager;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.data.redis.serializer.StringRedisSerializer;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Collections;

/**
 * Redis限流核心处理逻辑。在切面中进行处理
 *
 * @author biaoy
 * @since 2020/12/22
 */
@Slf4j
@Aspect
@Component
public class RateLimitAspectHandler {

    @Resource
    private RedisTemplate<String, String> redisTemplate;

    @Resource
    private RedisScript<String> redisScript;

    @Resource
    private StringRedisSerializer redisSerializer;

    @Resource
    private LimitKeyGenerator redisKeyGenerator;

    @Around("execution(public * *(..)) && @annotation(com.augrain.microframework.limit.annotation.RateLimiter)")
    public Object interceptor(ProceedingJoinPoint pjp) throws Throwable {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        RateLimiter redisLimit = method.getAnnotation(RateLimiter.class);
        String prefix = redisLimit.prefix();

        RateRule[] rules = redisLimit.rules();
        if (rules.length >= 1) {
            String key = redisKeyGenerator.generate(prefix, redisLimit.limitType(), pjp);
            boolean acquire = tryAcquireMultiRule(key, rules);
            if (acquire) {
                return pjp.proceed();
            } else {
                return CircuitBreakerManager.getHandler(redisLimit.breaker()).handle();
            }
        }
        return pjp.proceed();
    }

    private Boolean tryAcquireMultiRule(String key, RateRule[] rules) {
        // Lua 脚本中的 args
        Object[] args = new Object[rules.length * 2];
        for (int i = 0; i < rules.length; i++) {
            args[i * 2] = String.valueOf(rules[i].limit());
            args[i * 2 + 1] = String.valueOf(rules[i].timeUnit().toSeconds(rules[i].timeDuration()));
        }

        String result = redisTemplate.execute(redisScript, redisSerializer, redisSerializer,
                Collections.singletonList(key), args);
        log.info("结果：{}", result);
        return Boolean.valueOf(result);
    }
}
